package main

import (
	"context"
	"flag"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"github.com/jmoiron/sqlx"
	"github.com/panjf2000/ants/v2"
	"github.com/zeromicro/go-zero/core/conf"
	"math/rand"
	"strconv"
	"time"
)

var configFile = flag.String("f", "td-mysql.yaml", "the config file")

func main() {
	// 读取配置文件
	var c Config
	conf.MustLoad(*configFile, &c)

	fmt.Printf("%#v\n", c)

	svcCtx := NewServiceContext(c)

	var err error
	// 连接mysql
	svcCtx.Db, err = c.connectDB()
	if err != nil {
		// 可能不存在数据库  创建数据库

		err = c.createDatabase()
		if err != nil {
			panic(err)
		}

		// 重新连接mysql
		svcCtx.Db, err = c.connectDB()
		if err != nil {
			panic(err)
		}
	}

	fmt.Println("mysql连接成功")

	if c.AllLimit > 0 {
		svcCtx.DataAntsPool, _ = ants.NewPoolWithFunc(c.AllLimit, func(req interface{}) {
			data, _ := req.(int)
			svcCtx.Idata(data)
		})
	}

	// 创建数据表
	for i := 0; i < c.Limit; i++ {

		ctx, _ := context.WithTimeout(context.Background(), 60*time.Second)
		// 判断数据库表 是否存在
		data := fmt.Sprintf("SHOW CREATE TABLE d%d", i)
		_, err = svcCtx.Db.ExecContext(ctx, data)
		if err == nil {
			continue
		}

		name := fmt.Sprintf("CREATE TABLE d%d (ts BIGINT NOT NULL,value DOUBLE,PRIMARY KEY (ts)) ENGINE=ROCKSDB", i)

		_, err = svcCtx.Db.ExecContext(ctx, name)
		if err != nil {
			fmt.Println(err.Error())

		}
	}

	// 插入数据
	zero := 1

	for {
		for i := 0; i < c.Limit; i++ {
			if svcCtx.Config.AllLimit > 0 {
				_ = svcCtx.DataAntsPool.Invoke(i)
			} else {
				go svcCtx.Idata(i)
			}
		}

		fmt.Println(fmt.Sprintf("并发数%d已发送完第%d次", c.Limit, zero))
		time.Sleep(time.Second * time.Duration(c.TimeNext))

		// 计数
		zero++
	}

}

func (s ServiceContext) Idata(i int) {

	numberRandom := RandomDecimal("0", 300)

	ctxA, _ := context.WithTimeout(context.Background(), 60*time.Second)

	cc := fmt.Sprintf("INSERT INTO d%d (ts, value) VALUES ('%v', '%g')", i, time.Now().UnixMilli(), numberRandom)

	_, err := s.Db.ExecContext(ctxA, cc)
	if err != nil {

		fmt.Println(err.Error())
	}

}

func RandomDecimal(bit string, multiple float64) float64 {
	bit = "%." + bit + "f"
	data := rand.Float64() * multiple
	data, _ = strconv.ParseFloat(fmt.Sprintf(bit, data), 64)
	return data
}

func (c Config) connectDB() (*sqlx.DB, error) {
	// 连接到一个已存在的数据库，例如mysql数据库
	dataSourceName := fmt.Sprintf("%s:%s@tcp(%s)/%s", c.MysqlConfig.User, c.MysqlConfig.Password, c.MysqlConfig.Host, c.MysqlConfig.DbName)
	db, err := sqlx.Connect("mysql", dataSourceName)
	if err != nil {
		return nil, err
	}

	if err = db.Ping(); err != nil {
		return nil, err
	}
	return db, nil
}

func (c Config) createDatabase() error {
	// 使用root用户连接到MySQL服务器（通常是localhost或127.0.0.1）
	dataSourceName := fmt.Sprintf("%s:%s@tcp(%s)/", c.MysqlConfig.User, c.MysqlConfig.Password, c.MysqlConfig.Host)
	db, err := sqlx.Connect("mysql", dataSourceName)
	if err != nil {
		return err
	}

	defer db.Close() // 确保关闭连接

	if err = db.Ping(); err != nil {
		return err
	}

	// 创建新数据库的SQL命令
	_, err = db.Exec(fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s", c.MysqlConfig.DbName))
	if err != nil {
		return err
	}

	return nil
}
